﻿using Ndns.Packet.Abstractions;
using Ndns.Packet.Parser;
using System;
using System.Collections.Generic;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;

namespace Ndns.Net
{
    
    public class DnsTcpClient : IDisposable
    {
        private IPEndPoint _serverEP;
        private Socket _socket;
        private NetworkStream _ns;

        private volatile int _isBusy = 0;
        private IEnumerator<ushort> _sequence;

        private bool _useRandomTransactionId;
        public bool UseRandomTransactionId
        {
            get => _useRandomTransactionId;
            set
            {
                if (_useRandomTransactionId != value)
                {
                    _sequence.Dispose();
                    if (value)
                        _sequence = TransactionIdGenerator.GetRandomTransactionId().GetEnumerator();
                    else
                        _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
                    _useRandomTransactionId = value;
                }
            }
        }
        public IPEndPoint ServerEndPoint => _serverEP;

        public bool EnableCompression { get; set; } = true;
        
        private DnsPacketParser _packetParser;

        public DnsTcpClient()
        {
            _useRandomTransactionId = false;
            _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
            _packetParser = new DnsPacketParser();
        }

        public DnsTcpClient(DnsPacketParser packetParser)
        {
            _useRandomTransactionId = false;
            _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
            _packetParser = packetParser ?? new DnsPacketParser();
        }

        public Task ConnectAsync(IPEndPoint serverEP)
        {
            _ns?.Dispose();
            _socket?.Dispose();
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpClient));

            _serverEP = serverEP ?? throw new ArgumentNullException(nameof(serverEP));
            InternalInitSocket();
            return InternalConnect();
        }
        private void InternalInitSocket()
        {
            _socket = new Socket(_serverEP.AddressFamily, SocketType.Stream, ProtocolType.Tcp);
        }
        private async Task InternalConnect()
        {
            if (!_socket.Connected)
                await _socket.ConnectAsync(_serverEP).ConfigureAwait(false);
            _ns = new NetworkStream(_socket, false);
        }
        private async Task<IDisposable> InternalReconnect()
        {
            var socket = new Socket(_socket.AddressFamily, _socket.SocketType, _socket.ProtocolType);
            _socket.Dispose();
            _socket = socket;
            await _socket.ConnectAsync(_serverEP).ConfigureAwait(false);
            _ns = new NetworkStream(_socket, false);
            return _socket;
        }


        public async Task<DnsPacket> SendAsync(DnsQuery query, CancellationToken cancellationToken = default(CancellationToken))
        {
            if (query == null)
                throw new ArgumentNullException(nameof(query));
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpClient));
            if (_socket == null)
                throw new InvalidOperationException();
            if (_isBusy == 1)
                throw new InvalidOperationException();
            if (1 == Interlocked.Exchange(ref _isBusy, 1))
                throw new InvalidOperationException();

            _sequence.MoveNext();
            DnsPacket request = new DnsPacket(_sequence.Current, DnsPacketFlag.RecursionDesired);
            request.Queries.Add(query);

            try
            {
                return await InternalSendAsync(request, cancellationToken).ConfigureAwait(false);
            }
            finally
            {
                _isBusy = 0;
            }
        }

        public async Task<DnsPacket> SendAsync(DnsPacket request, CancellationToken cancellationToken = default(CancellationToken))
        {
            if (request == null)
                throw new ArgumentNullException(nameof(request));
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpClient));
            if (_socket == null)
                throw new InvalidOperationException();
            if (_isBusy == 1)
                throw new InvalidOperationException();
            if (1 == Interlocked.Exchange(ref _isBusy, 1))
                throw new InvalidOperationException();

            try
            {
                return await InternalSendAsync(request, cancellationToken).ConfigureAwait(false);
            }
            finally
            {
                _isBusy = 0;
            }
        }

        private async Task<DnsPacket> InternalSendAsync(DnsPacket request, CancellationToken cancellationToken)
        {
            using (cancellationToken.Register(() => _socket.Dispose()))
            {
                try
                {
                    //确保连接
                    await InternalConnect().ConfigureAwait(false);
                    using (var bs = new MemoryStream(512)) // Used for buffer
                    using (var ps = new MemoryStream(512)) // Used for packet
                    {
                        //发送数据
                        ps.SetLength(2);
                        ps.Seek(2, SeekOrigin.Begin);
                        int count = EnableCompression ? request.CopyCompressedTo(ps) : request.CopyTo(ps);
                        ps.Seek(0, SeekOrigin.Begin);
                        ps.WriteNetworkUInt16((ushort)count);
                        ps.Seek(0, SeekOrigin.Begin);
                        await ps.CopyToAsync(_ns).ConfigureAwait(false);
                        //接收数据
                        await _ns.CopyBytesToAsync(bs, 2).ConfigureAwait(false);
                        bs.Seek(0, SeekOrigin.Begin);
                        int expectedLength = bs.ReadNetworkUInt16();
                        if (expectedLength == -1)
                        {
                            //接收失败，可能是连接被断开，重新发送
                            var conn = await InternalReconnect().ConfigureAwait(false);
                            using (cancellationToken.Register(conn.Dispose))
                            {
                                //重新发送
                                ps.Seek(0, SeekOrigin.Begin);
                                await ps.CopyToAsync(_ns).ConfigureAwait(false);
                                //重新接收
                                bs.Seek(0, SeekOrigin.Begin);
                                await _ns.CopyBytesToAsync(bs, 2).ConfigureAwait(false);
                                bs.Seek(0, SeekOrigin.Begin);
                                expectedLength = bs.ReadNetworkUInt16();
                                //如果还是接收失败
                                if (expectedLength == -1)
                                    throw new IOException();
                            }
                        }
                        await _ns.CopyBytesToAsync(bs, expectedLength).ConfigureAwait(false);
                        bs.Seek(2, SeekOrigin.Begin);
                        if (!_packetParser.TryParse(bs, out var packet, out var c) && !(packet?.Truncated ?? false))
                            throw new FormatException();
                        return packet;
                    }
                }
                catch (ObjectDisposedException)
                {
                    InternalInitSocket();
                    cancellationToken.ThrowIfCancellationRequested();
                    throw;
                }
            }

        }

        public static async Task<DnsPacket> SendAsync(DnsPacket request, IPEndPoint server, CancellationToken cancellationToken = default(CancellationToken))
        {
            var _parser = new DnsPacketParser();

            //准备Socket
            using (var tcpSocket = new Socket(server.AddressFamily, SocketType.Stream, ProtocolType.Tcp))
            using (cancellationToken.Register(() => tcpSocket.Dispose()))
            {
                try
                {
                    DnsPacket packet;
                    //连接服务器
                    await tcpSocket.ConnectAsync(server.Address, server.Port).ConfigureAwait(false);
                    using (var ns = new NetworkStream(tcpSocket, false))
                    using (var ms = new MemoryStream(512))
                    {
                        //发送数据
                        ms.SetLength(2);
                        ms.Seek(2, SeekOrigin.Begin);
                        int count = request.CopyCompressedTo(ms);
                        ms.Seek(0, SeekOrigin.Begin);
                        ms.WriteNetworkUInt16((ushort)count);
                        ms.Seek(0, SeekOrigin.Begin);
                        await ms.CopyToAsync(ns).ConfigureAwait(false);
                        //接收数据
                        ms.Seek(0, SeekOrigin.Begin);
                        ms.SetLength(0);
                        await ns.CopyBytesToAsync(ms, 2).ConfigureAwait(false);
                        ms.Seek(0, SeekOrigin.Begin);
                        int expectedLength = ms.ReadNetworkUInt16();
                        await ns.CopyBytesToAsync(ms, expectedLength).ConfigureAwait(false);
                        ms.Seek(2, SeekOrigin.Begin);
                        if (!_parser.TryParse(ms, out packet, out var c) && !(packet?.Truncated ?? false))
                            throw new FormatException();
                    }
                    //关闭连接
                    tcpSocket.Shutdown(SocketShutdown.Both);
                    //返回
                    return packet;
                }
                catch (ObjectDisposedException)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    throw;
                }
            }
        }
        

        #region IDisposable Support
        private bool disposedValue = false; // 要检测冗余调用

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    // 释放托管状态(托管对象)。
                    _ns?.Dispose();
                    _socket?.Dispose();
                    _sequence?.Dispose();
                }

                // 释放未托管的资源(未托管的对象)并在以下内容中替代终结器。
                // 将大型字段设置为 null。
                _ns = null;
                _socket = null;
                _sequence = null;

                disposedValue = true;
            }
        }
        
        // 添加此代码以正确实现可处置模式。
        public void Dispose()
        {
            // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
            Dispose(true);
        }
        #endregion

    }
}
